/**
 * @file scatter_reduce.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "scatter_reduce_tiling.h"
#include "kernel_operator.h"
#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

// Define reduction operations
enum class ReduceOp
{
    SUM,
    PROD,
    MEAN,
    AMAX,
    AMIN
};

template <typename T_VAR, typename T_INDICES>
class ScatterReduce
{
public:
    __aicore__ inline ScatterReduce() {}
    __aicore__ inline void Init(GM_ADDR self, GM_ADDR index, GM_ADDR src, GM_ADDR workspace,
                                uint32_t blockNum,
                                ReduceType reduce,
                                bool include_self,
                                uint32_t dim,
                                uint32_t blockLength,
                                uint32_t lastBlockLength,
                                uint32_t prevDimSize,
                                uint32_t selfDimSize,
                                uint32_t indexDimSize,
                                uint32_t postDimSize)
    {
        this->blockNum = blockNum;
        this->reduce = reduce;
        this->include_self = include_self;
        this->dim = dim;
        this->blockLength = blockLength;
        this->lastBlockLength = lastBlockLength;
        this->prevDimSize = prevDimSize;
        this->selfDimSize = selfDimSize;
        this->indexDimSize = indexDimSize;
        this->postDimSize = postDimSize;

        auto blockIdx = AscendC::GetBlockIdx();
        this->thisBlockLength = (blockIdx == this->blockNum - 1) ? this->lastBlockLength : this->blockLength;

        this->selfGm.SetGlobalBuffer((__gm__ T_VAR *)self + blockIdx * this->blockLength * this->selfDimSize * this->postDimSize,
                                     this->thisBlockLength * this->selfDimSize * this->postDimSize);
        this->indexGm.SetGlobalBuffer((__gm__ T_INDICES *)index + blockIdx * this->blockLength * this->indexDimSize * this->postDimSize,
                                      this->thisBlockLength * this->indexDimSize * this->postDimSize);
        this->srcGm.SetGlobalBuffer((__gm__ T_VAR *)src + blockIdx * this->blockLength * this->indexDimSize * this->postDimSize,
                                    this->thisBlockLength * this->indexDimSize * this->postDimSize);
        this->counterGm.SetGlobalBuffer((__gm__ T_INDICES *)workspace + blockIdx * this->blockLength * this->selfDimSize * this->postDimSize,
                                        this->thisBlockLength * this->selfDimSize * this->postDimSize);
                                    
    }

    __aicore__ inline void ProcessSum()
    {
        for (uint32_t i = 0; i < this->thisBlockLength; i++)
        {
            for (uint32_t j = 0; j < this->indexDimSize; j++)
            {
                for (uint32_t k = 0; k < this->postDimSize; k++)
                {
                    uint32_t index = this->indexGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k);
                    uint32_t targetIdx = (i * this->selfDimSize + index) * this->postDimSize + k;
                    uint32_t currentCount = this->counterGm.GetValue(targetIdx);
                    
                    if (!this->include_self && currentCount == 0)
                    {
                        // First update and include_self is false - overwrite the value
                        this->selfGm.SetValue(targetIdx,
                                            this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k));
                        // Increment counter
                        this->counterGm.SetValue(targetIdx, currentCount + 1);
                    }
                    else
                    {
                        // Add the value for subsequent updates or when include_self is true
                        this->selfGm.SetValue(targetIdx,
                                            this->selfGm.GetValue(targetIdx) +
                                            this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k));
                    }
                    
                }
            }
        }
    }

    __aicore__ inline void ProcessProd()
    {
        for (uint32_t i = 0; i < this->thisBlockLength; i++)
        {
            for (uint32_t j = 0; j < this->indexDimSize; j++)
            {
                for (uint32_t k = 0; k < this->postDimSize; k++)
                {
                    uint32_t index = this->indexGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k);
                    uint32_t targetIdx = (i * this->selfDimSize + index) * this->postDimSize + k;
                    uint32_t currentCount = this->counterGm.GetValue(targetIdx);
                    
                    if (!this->include_self && currentCount == 0)
                    {
                        // First update and include_self is false - overwrite the value
                        this->selfGm.SetValue(targetIdx,
                                            this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k));
                        // Increment counter
                        this->counterGm.SetValue(targetIdx, currentCount + 1);
                    }
                    else
                    {
                        // Multiply the value for subsequent updates or when include_self is true
                        this->selfGm.SetValue(targetIdx,
                                            this->selfGm.GetValue(targetIdx) *
                                            this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k));
                    }
                }
            }
        }
    }

    __aicore__ inline void ProcessMean()
    {
        // Initialize counters for include_self=true case
        if (this->include_self) {
            for (uint32_t i = 0; i < this->thisBlockLength; i++) {
                for (uint32_t index = 0; index < this->selfDimSize; index++) {
                    for (uint32_t k = 0; k < this->postDimSize; k++) {
                        uint32_t targetIdx = (i * this->selfDimSize + index) * this->postDimSize + k;
                        this->counterGm.SetValue(targetIdx, 1); // Start with count 1 for include_self
                    }
                }
            }
        }
        
        for (uint32_t i = 0; i < this->thisBlockLength; i++)
        {
            for (uint32_t j = 0; j < this->indexDimSize; j++)
            {
                for (uint32_t k = 0; k < this->postDimSize; k++)
                {
                    uint32_t index = this->indexGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k);
                    uint32_t targetIdx = (i * this->selfDimSize + index) * this->postDimSize + k;
                    uint32_t currentCount = this->counterGm.GetValue(targetIdx);
                    
                    if (!this->include_self && currentCount == 0)
                    {
                        // First update and include_self is false - overwrite the value
                        this->selfGm.SetValue(targetIdx,
                                            this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k));
                        // Set counter to 1
                        this->counterGm.SetValue(targetIdx, 1);
                    }
                    else
                    {
                        // Add the value for subsequent updates or when include_self is true
                        this->selfGm.SetValue(targetIdx,
                                            this->selfGm.GetValue(targetIdx) +
                                            this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k));
                        // Increment counter - needed for later averaging
                        this->counterGm.SetValue(targetIdx, currentCount + 1);
                    }
                }
            }
        }
        
        // After all values are accumulated, calculate the mean by dividing by count
        for (uint32_t i = 0; i < this->thisBlockLength; i++)
        {
            for (uint32_t index = 0; index < this->selfDimSize; index++)
            {
                for (uint32_t k = 0; k < this->postDimSize; k++)
                {
                    uint32_t targetIdx = (i * this->selfDimSize + index) * this->postDimSize + k;
                    uint32_t count = this->counterGm.GetValue(targetIdx);
                    if (count > 0) {
                        this->selfGm.SetValue(targetIdx, this->selfGm.GetValue(targetIdx) / count);
                    }
                }
            }
        }
    }

    __aicore__ inline void ProcessAmax()
    {
        for (uint32_t i = 0; i < this->thisBlockLength; i++)
        {
            for (uint32_t j = 0; j < this->indexDimSize; j++)
            {
                for (uint32_t k = 0; k < this->postDimSize; k++)
                {
                    uint32_t index = this->indexGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k);
                    uint32_t targetIdx = (i * this->selfDimSize + index) * this->postDimSize + k;
                    uint32_t currentCount = this->counterGm.GetValue(targetIdx);
                    
                    if (!this->include_self && currentCount == 0)
                    {
                        // First update and include_self is false - overwrite the value
                        this->selfGm.SetValue(targetIdx,
                                            this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k));
                        // Increment counter
                        this->counterGm.SetValue(targetIdx, currentCount + 1);
                    }
                    else
                    {
                        // Take maximum value
                        this->selfGm.SetValue(targetIdx,
                                            MAX(this->selfGm.GetValue(targetIdx),
                                                this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k)));
                    }
                }
            }
        }
    }

    __aicore__ inline void ProcessAmin()
    {
        for (uint32_t i = 0; i < this->thisBlockLength; i++)
        {
            for (uint32_t j = 0; j < this->indexDimSize; j++)
            {
                for (uint32_t k = 0; k < this->postDimSize; k++)
                {
                    uint32_t index = this->indexGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k);
                    uint32_t targetIdx = (i * this->selfDimSize + index) * this->postDimSize + k;
                    uint32_t currentCount = this->counterGm.GetValue(targetIdx);
                    
                    if (!this->include_self && currentCount == 0)
                    {
                        // First update and include_self is false - overwrite the value
                        this->selfGm.SetValue(targetIdx,
                                            this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k));
                        // Increment counter
                        this->counterGm.SetValue(targetIdx, currentCount + 1);
                    }
                    else
                    {
                        // Take minimum value
                        this->selfGm.SetValue(targetIdx,
                                            MIN(this->selfGm.GetValue(targetIdx),
                                                this->srcGm.GetValue((i * this->indexDimSize + j) * this->postDimSize + k)));
                    }
                }
            }
        }
    }
    
private:
    AscendC::GlobalTensor<T_VAR> selfGm;
    AscendC::GlobalTensor<T_INDICES> indexGm;
    AscendC::GlobalTensor<T_VAR> srcGm;
    AscendC::GlobalTensor<T_INDICES> counterGm;
    uint32_t blockNum;
    uint32_t thisBlockLength;
    ReduceType reduce;
    bool include_self;
    uint32_t dim;
    uint32_t blockLength;
    uint32_t lastBlockLength;
    uint32_t totalLength;
    uint32_t prevDimSize;
    uint32_t selfDimSize;
    uint32_t indexDimSize;
    uint32_t postDimSize;
};

extern "C" __global__ __aicore__ void scatter_reduce(GM_ADDR self, GM_ADDR index, GM_ADDR src, GM_ADDR currentWorkspace, ScatterReduceTilingData tiling)
{
    // Implementation for scatter reduce operation
    ScatterReduce<DTYPE_SELF, DTYPE_INDEX> op;
    op.Init(self, index, src, currentWorkspace,
            tiling.blockNum,
            tiling.reduce,
            tiling.include_self,
            tiling.dim,
            tiling.blockLength,
            tiling.lastBlockLength,
            tiling.prevDimSize,
            tiling.selfDimSize,
            tiling.indexDimSize,
            tiling.postDimSize);
    switch (tiling.reduce)
    {
    case ReduceType::SUM:
        op.ProcessSum();
        break;
    case ReduceType::PROD:
        op.ProcessProd();
        break;
    case ReduceType::MEAN:
        op.ProcessMean();
        break;
    case ReduceType::AMAX:
        op.ProcessAmax();
        break;
    case ReduceType::AMIN:
        op.ProcessAmin();
        break;
    }
}